#!/usr/bin/env python2.5
#coding: utf8

#Author: Gang Li
#Date: 2010-04-10 18:20

import time, select, socket
from threading import Thread as _t, RLock

from core.utils import *
from core.transport import Server, CmdFilter, HeartBeat, CmdDispatcher
from core.fs import FSRoot
from core.task import Task, SubTask

import logging
logging.basicConfig(format="%(asctime)s %(levelname)s %(message)s", level=logging.INFO)
log = logging.getLogger("controller")

class Controller(CmdFilter):
    """
    free
    alloc
    wait_task
    """
    hb = HeartBeat()
    def __init__(self, binding_ip, port):
        self.server = Server(binding_ip, port, self)
        self.server.add_cmd_observer(self)
        self.server.add_cmd_observer(self.hb)
        self.cid = 0
        self.cid_lock = RLock()
        self.task_lock = RLock()

        self.tasks = {}
        self.nodes = {}

        #recevied task
        self.waiting_tasks = []

        self.fs = FSRoot.open(mode="meta", handler=self)
        self.fs.get_root().mkdirs("/x/y/z/u/w/i")

    def allocate_free_nodes(self, size, gid):
        al = []

        if size > len(self.nodes):
            return

        self.task_lock.acquire()
        try:
            for n in self.nodes:
                node_state = self.nodes[n].get("state")
                if node_state == 'free':
                    al.append(n)

                    if len(al) == size:
                        for a in al:
                            self.nodes[a]['state'] = 'alloc'
                            log.info("send alloc node cmd for %s" % gid)
                            self.server.write(a, TASK_NODE_ALLOC, gid)
                        return al
        finally:
            self.task_lock.release()

    def task_loop_r(self, task):
        task.reduces = []
        for i in range(task.r):
            st = SubTask(task.id, task.name, type="reduce", res=task.res)
            st.id = hash(st)

            self.server.write(task.nodes[task.m + i], START_SUB_TASK, st)
            task.reduces.append(st.id)
            log.info("start reduce func")
            yield st

    def task_loop(self):
        for task in self.waiting_tasks:
            nodes = self.allocate_free_nodes(task.m + task.r, task.id)

            if nodes:
                path = task.dir

                task.nodes = nodes
                self.tasks[task.id] = task

                #sub map task list
                task.maps = []
                for i in range(task.m):
                    st = SubTask(task.id, task.name, type="map", res=task.res)
                    st.id = hash(st)

                    self.server.write(nodes[i], START_SUB_TASK, st)
                    task.maps.append(st.id)

                task.start_reduce_task = self.task_loop_r(task)
                self.waiting_tasks.remove(task)
            else:
                log.info("allocate task resources error, need:%s current: %s" % (task.m + task.r, len(self.nodes)))

    def connected(self, sock): 
        log.info("connnected, %s", sock.getsockname())
        self.nodes[sock] = {"state":"wait"}
        def sender(id):
            log.debug("try send hb msg, %s %s", id, sock)
            self.server.write(sock, id)
        def on_failure(skt):
            log.warn("hb failure %s", sock.getsockname())

        self.hb.add_observed(sock, sender, on_failure)
        #try to allocate task resource
        self.task_loop()

    def disconnected(self, sock): 
        log.info("disconnnected, %s", sock.getsockname())
        del self.nodes[sock]

    def on_cmd(self, msg):
        id, data, skt = msg
        try:
            if id == NODE_REG:
                log.info("got node reg for %s %s", skt.getsockname(), data)
                self.nodes[skt]["state"] = (data == "tasknode") and "free"  or "other"
            elif id.startswith(FS_CMD):
                if self.fs: self.fs.on_cmd(id, len, data)
            elif id == START_TASK:
                #user data
                task = data
                log.info("start task")
                task.id = hash(data)
                self.waiting_tasks.append(data)

                log.info("start task ok")
            elif id == TASK_NODE_ALLOC_RSP:
                self.nodes[skt]['state'] = 'wait_task'
                gid = int(data)
                log.info("recv alloc node cmd rsp for %s" % gid)
            elif id == START_SUB_TASK_OK:
                self.nodes[skt]['state'] = 'running_task'
            elif id == SUB_TASK_DONE:
                self.nodes[skt]['state'] = 'free'
                st = data

                if hasattr(st, "error"):
                    log.error(skt.getsockname(), st.error)

                task = self.tasks[st.pid]

                if st.type == 'map':
                    task.maps.remove(st.id)

                    if not task.maps:
                        log.info("map done!!!")

                    for st in task.start_reduce_task:
                        log.info("start reduce!!!")
                elif st.type == 'reduce':
                    task.reduces.remove(st.id)

                    if not task.reduces:
                        log.info("reduce done!!!, task done.")        
                        task.nodes = []

            #try to allocate task resource
            self.task_loop()
        except BaseException, msg:
            log.warn("process cmd error:  %s", msg)
            import traceback
            traceback.print_exc()

    def get_cid(self):
        self.cid_lock.acquire()
        try:
            self.cid += 1
            return self.cid
        finally:
            self.cid_lock.release()
    def reg_file(self, param):
        self.server.write_all(REG_FILE, param)
    def reg_dir(self, param):
        self.server.write_all(REG_DIR , param)
